package com.ciccwm.log.kafka;

import com.ciccwm.log.constant.LogConstant;
import com.ciccwm.log.util.TraceIdUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.ConsumerInterceptor;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.header.Header;
import org.springframework.util.StringUtils;

import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * Kafka消费者链路追踪拦截器.
 * 用于在消费Kafka消息时，从消息头中提取追踪ID并设置到当前线程上下文中.
 * 支持在分布式系统中追踪消息的完整处理链路.
 */
@Slf4j
public class TraceIdConsumerInterceptor<K, V> implements ConsumerInterceptor<K, V> {
    /**
     * 在消费消息时处理追踪ID.
     * 从消息头中提取追踪ID，并设置到当前线程上下文中.
     * 使用executeWithTraceId确保每个消息都在正确的追踪上下文中处理.
     *
     * @param records 消费者记录
     * @return 处理后的消费者记录
     */
    @Override
    public ConsumerRecords<K, V> onConsume(ConsumerRecords<K, V> records) {
        records.forEach(record -> {
            Header traceIdHeader = record.headers().lastHeader(LogConstant.TRACE_ID_KEY);
            if (traceIdHeader != null) {
                String traceId = new String(traceIdHeader.value(), StandardCharsets.UTF_8);
                if (StringUtils.hasText(traceId)) {
                    log.debug("Found trace ID in Kafka consumer record: {}", traceId);
                    // todo 什么是否设置traceId
                    TraceIdUtil.executeWithTraceId(traceId,
                            () -> log.debug("Processing message with trace ID: {}", traceId));
                }
            }
        });
        return records;
    }

    /**
     * 在提交偏移量时处理.
     * 当前实现不执行任何操作.
     *
     * @param offsets 主题分区和偏移量元数据的映射
     */
    @Override
    public void onCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {
        // Do nothing
    }

    /**
     * 关闭拦截器.
     * 当前实现不执行任何操作.
     */
    @Override
    public void close() {
        // Do nothing
    }

    /**
     * 配置拦截器.
     * 当前实现不执行任何操作.
     *
     * @param configs 配置参数
     */
    @Override
    public void configure(Map<String, ?> configs) {
        // Do nothing
    }
} 
